#!/usr/bin/env python
# -*- coding = utf-8 -*-
# @Time      : 2021/1/18 17:01
# @Author    : Reanon
# @File      : ast_to_graph
# @Email     : dafo360@gmail.com
# @Software  : PyCharm


# 保存节点
import json

import numpy as np

# 定义节点
contracts_mapping = {
    'ContractDefinition': ['baseContracts', 'subNodes'],
    'FunctionDefinition': ['parameters', 'returnParameters', 'body'],
    'StateVariableDeclaration': ['variables'],
    'ParameterList': ['parameters'],
    'parameters': ['Parameter', 'typeName'],
    'Block': ['statements'],
    'ExpressionStatement': ['expression']
}
contracts = {
    'ContractDefinition': ['baseContracts', 'subNodes'],
    'FunctionDefinition': ['parameters', 'returnParameters', 'body'],
    'StateVariableDeclaration': ['variables'],
    'ParameterList': ['parameters'],
    'parameters': ['Parameter', 'typeName'],
    'ExpressionStatement': ['expression']
}
nodes_value = {
    'AST': 0,
    'ContractDefinition': 1,
    'FunctionDefinition': 2,
    'StateVariableDeclaration': 3,
    'ParameterList': 4,
    'parameters': 5,
    'ExpressionStatement': 6,
}

# 记录全局节点总数
global counts
counts = 0


def load():
    contract_name = "input"
    # 加载数据
    filename = "../ast_json/" + contract_name + ".json"
    with open(filename, 'r') as file:
        # data 是字典类型
        ast_data = json.load(file)
        # 检查是否为 字典 类型
        print(isinstance(ast_data, dict))
    return ast_data


def dfs_nodes_name(prev_node, current_node, nodes, edge_index):
    """
    深度优先搜索，生成图的节点名组成的图
    :param prev_node:
    :param current_node:
    :param nodes:
    :param edge_index:
    :param counts:
    :return:
    """
    global counts
    pre_node_type = prev_node['node_type']

    # 如果当前节点 有类型
    if 'type' in current_node:
        # 获取当前节点类型
        current_node_type = current_node['type']
    else:
        return

    # 如果当前节点类型 属于图节点类型，则添加该节点
    if current_node_type in contracts:
        # 增加图节点的总数量
        counts += 1
        # 增加一个图节点
        nodes.append(current_node_type + "_" + str(counts))

        # 在上个图节点 与 该图节点之间 增加一条边
        edge_index.append([pre_node_type + "_" + str(prev_node['counts']),
                           current_node_type + "_" + str(counts)])

        # 更新 节点索引，只需要传节点类型和节点数即可
        prev_node['node_type'] = current_node_type
        prev_node['counts'] = counts

    # 当前节点是否在 contracts表 中有记录
    if current_node_type in contracts_mapping:
        for child_node_type in contracts_mapping[current_node_type]:
            # 如果当前子节点是个 列表，如subNodes,returnParameters等
            if isinstance(current_node[child_node_type], list):
                for subNode in current_node[child_node_type]:
                    dfs_nodes_value(prev_node, subNode, nodes, edge_index)
            else:
                # 如果 子节点 是一个节点类型，则对子节点进行遍历
                child_node = current_node[child_node_type]
                dfs_nodes_value(prev_node, child_node, nodes, edge_index)


def ast_graph_name():
    """
    从图生成树
    :return:
    """
    # 加载数据
    data = load()
    # 生成带有图数据的边
    nodes = []
    edge_index = []
    prev_node = {'node_type': 'AST', 'counts': 0}
    nodes = ['AST_0']
    for current_node in data["children"]:
        if 'type' in current_node:
            current_node['node_type'] = current_node['type']
            dfs_nodes_value(prev_node, current_node, nodes, edge_index)

    return [nodes, edge_index]


def dfs_nodes_value(prev_node, current_node, nodes, edge_index):
    """
    深度优先搜索，生成图的节点信息
    :param prev_node:
    :param current_node:
    :param nodes:
    :param edge_index:
    :param counts:
    :return:
    """
    global counts
    pre_node_type = prev_node['node_type']

    # 如果当前节点 有类型
    if 'type' in current_node:
        # 获取当前节点类型
        current_node_type = current_node['type']
    else:
        return

    # 如果当前节点类型 属于图节点类型，则添加该节点
    if current_node_type in contracts:
        # 增加图节点的总数量
        counts += 1
        # 增加一个图节点,对应节点的值
        value = nodes_value[current_node_type]
        nodes.append(value)

        # 在上个图节点 与 该图节点之间 增加一条边
        edge_index.append([prev_node['counts'],
                           counts])

        # 更新 节点索引，只需要传节点类型和节点数即可
        prev_node['node_type'] = current_node_type
        prev_node['counts'] = counts

    # 当前节点是否在 contracts表 中有记录
    if current_node_type in contracts_mapping:
        for child_node_type in contracts_mapping[current_node_type]:
            # 如果当前子节点是个 列表，如subNodes,returnParameters等
            if isinstance(current_node[child_node_type], list):
                for subNode in current_node[child_node_type]:
                    dfs_nodes_value(prev_node, subNode, nodes, edge_index)
            else:
                # 如果 子节点 是一个节点类型，则对子节点进行遍历
                child_node = current_node[child_node_type]
                dfs_nodes_value(prev_node, child_node, nodes, edge_index)


def ast_graph_value():
    """
    从图生成树
    :return:
    """
    # 加载数据
    data = load()
    # 生成带有图数据的边
    nodes = []
    edge_index = []
    prev_node = {'node_type': 'AST', 'counts': 0}
    nodes = [0]
    for current_node in data["children"]:
        if 'type' in current_node:
            current_node['node_type'] = current_node['type']
            dfs_nodes_value(prev_node, current_node, nodes, edge_index)

    return [nodes, edge_index]


if __name__ == '__main__':
    nodes_list, edge_index_list = ast_graph_value()
    edge_index1 = np.array(edge_index_list)
    nodes1 = np.array(nodes_list)
    # 打印节点和边信息
    print("AST的图节点信息：\n{} \nAST的边信息：\n{}"
          .format(nodes1, edge_index1))